#include "onnxruntime_cxx_api.h"
#include "cpu_provider_factory.h"
#include <opencv2/opencv.hpp>
#include <fstream>

std::string labels_txt_file = "/home/chyang/Documents/program/deploy/onnxruntime/onnxruntime_cpp/model/imagenet_classes.txt";
std::vector<std::string> readClassNames();
std::vector<std::string> readClassNames()
{
    std::vector<std::string> classNames;

    std::ifstream fp(labels_txt_file);
    if (!fp.is_open())
    {
        printf("could not open file...\n");
        exit(-1);
    }
    std::string name;
    while (!fp.eof())
    {
        std::getline(fp, name);
        if (name.length())
            classNames.push_back(name);
    }
    fp.close();
    return classNames;
}

int main(int argc, char **argv)
{
    std::vector<std::string> labels = readClassNames();
    cv::Mat image = cv::imread("/home/chyang/Documents/program/deploy/onnxruntime/onnxruntime_cpp/images/car_test.png");
    cv::imshow("输入图", image);
    // 创建InferSession, 查询支持硬件设备
    // GPU Mode, 0 - gpu device id
    std::string onnxpath = "/home/chyang/Documents/program/deploy/onnxruntime/onnxruntime_cpp/model/resnet18.onnx";
    // std::wstring modelPath = std::wstring(onnxpath.begin(), onnxpath.end());
    Ort::SessionOptions session_options;
    Ort::Env env = Ort::Env(ORT_LOGGING_LEVEL_ERROR, "resnet18-onnx");

    session_options.SetGraphOptimizationLevel(ORT_ENABLE_BASIC);
    std::cout << "onnxruntime inference try to use GPU Device" << std::endl;
    // OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0);
    
    OrtSessionOptionsAppendExecutionProvider_CPU(session_options, 0);
    Ort::Session session_(env, onnxpath.c_str(), session_options);

    // get input and output info
    int input_nodes_num = session_.GetInputCount();
    int output_nodes_num = session_.GetOutputCount();
    std::vector<std::string> input_node_names;
    std::vector<std::string> output_node_names;
    Ort::AllocatorWithDefaultOptions allocator;
    int input_h = 0;
    int input_w = 0;

    std::vector<std::string> available_providers = Ort::GetAvailableProviders();

    // query input data format
    for (int i = 0; i < input_nodes_num; i++)
    {
        auto input_name = session_.GetInputNameAllocated(i, allocator);
        input_node_names.push_back(input_name.get());
        auto inputShapeInfo = session_.GetInputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape();
        int ch = inputShapeInfo[1];
        input_h = inputShapeInfo[2];
        input_w = inputShapeInfo[3];
        std::cout << "input format: " << ch << "x" << input_h << "x" << input_w << std::endl;
    }

    // query output data format
    int num = 0;
    int nc = 0;
    for (int i = 0; i < output_nodes_num; i++)
    {
        auto output_name = session_.GetOutputNameAllocated(i, allocator);
        output_node_names.push_back(output_name.get());
        auto outShapeInfo = session_.GetOutputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape();
        num = outShapeInfo[0];
        nc = outShapeInfo[1];
        std::cout << "output format: " << num << "x" << nc << std::endl;
    }

    // set input image
    cv::Mat rgb, blob;
    // RGB order
    cv::cvtColor(image, rgb, cv::COLOR_BGR2RGB);
    cv::resize(rgb, blob, cv::Size(input_w, input_h));
    blob.convertTo(blob, CV_32F);
    blob = blob / 255.0;
    cv::subtract(blob, cv::Scalar(0.485, 0.456, 0.406), blob);
    cv::divide(blob, cv::Scalar(0.229, 0.224, 0.225), blob);

    cv::Mat timg = cv::dnn::blobFromImage(blob);
    std::cout << timg.size[0] << "x" << timg.size[1] << "x" << timg.size[2] << "x" << timg.size[3] << std::endl;
    size_t tpixels = input_h * input_w * 3;
    std::array<int64_t, 4> input_shape_info{1, 3, input_h, input_w};

    // set input data and inference
    auto allocator_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
    Ort::Value input_tensor_ = Ort::Value::CreateTensor<float>(allocator_info, timg.ptr<float>(), tpixels, input_shape_info.data(), input_shape_info.size());
    const std::array<const char *, 1> inputNames = {input_node_names[0].c_str()};
    const std::array<const char *, 1> outNames = {output_node_names[0].c_str()};

    std::vector<Ort::Value> ort_outputs;
    try
    {
        ort_outputs = session_.Run(Ort::RunOptions{nullptr}, inputNames.data(), &input_tensor_, 1, outNames.data(), outNames.size());
    }
    catch (std::exception e)
    {
        std::cout << e.what() << std::endl;
    }
    // 1x 1000
    const float *pdata = ort_outputs[0].GetTensorMutableData<float>();
    cv::Mat prob(num, nc, CV_32F, (float *)pdata);

    // post process
    cv::Point maxL, minL;
    double maxv, minv;
    cv::minMaxLoc(prob, &minv, &maxv, &minL, &maxL);
    int max_index = maxL.x;
    std::cout << "label id: " << max_index << std::endl;
    cv::putText(image, labels[max_index], cv::Point(50, 50), cv::FONT_HERSHEY_SIMPLEX, 1.0, cv::Scalar(0, 0, 255), 2, 8);
    cv::imshow("输入图像", image);
    cv::waitKey(0);

    // relase resource
    session_options.release();
    session_.release();
    return 0;
}